package cn.edu.hitsz.compiler.asm;

import cn.edu.hitsz.compiler.ir.*;
import cn.edu.hitsz.compiler.utils.FileUtils;

import java.util.*;


/**
 * TODO: 实验四: 实现汇编生成
 * <br>
 * 在编译器的整体框架中, 代码生成可以称作后端, 而前面的所有工作都可称为前端.
 * <br>
 * 在前端完成的所有工作中, 都是与目标平台无关的, 而后端的工作为将前端生成的目标平台无关信息
 * 根据目标平台生成汇编代码. 前后端的分离有利于实现编译器面向不同平台生成汇编代码. 由于前后
 * 端分离的原因, 有可能前端生成的中间代码并不符合目标平台的汇编代码特点. 具体到本项目你可以
 * 尝试加入一个方法将中间代码调整为更接近 risc-v 汇编的形式, 这样会有利于汇编代码的生成.
 * <br>
 * 为保证实现上的自由, 框架中并未对后端提供基建, 在具体实现时可自行设计相关数据结构.
 *
 * @see AssemblyGenerator#run() 代码生成与寄存器分配
 */

public class AssemblyGenerator {

    // 存储中间代码指令的列表
    List<Instruction> instructions;
    // 存储生成的汇编代码的列表
    List<String> assembly_instructions = new ArrayList<>();

    /**
     * 加载前端提供的中间代码
     *
     * @param originInstructions 前端提供的中间代码
     */
    public void loadIR(List<Instruction> originInstructions) {
        this.instructions = new LinkedList<>(originInstructions);

        // 处理每一条中间代码指令，将指令调整为更接近 RISC-V 汇编的形式
        for (ListIterator<Instruction> iterator = instructions.listIterator(); iterator.hasNext(); ) {

            Instruction instr = iterator.next();
            InstructionKind instr_kind = instr.getKind();

            // 处理二元运算指令
            if (instr_kind.isBinary()) {
                List<IRValue> operands = instr.getOperands();
                IRVariable target = instr.getResult();
                IRValue operand1 = operands.get(0);
                IRValue operand2 = operands.get(1);

                // 如果两个操作数都是立即数，直接计算结果并生成 MOV 指令
                if (operand1.isImmediate() && operand2.isImmediate()) {
                    IRImmediate imm_operand1 = (IRImmediate) operand1;
                    IRImmediate imm_operand2 = (IRImmediate) operand2;
                    int result = switch (instr_kind) {
                        case ADD -> imm_operand1.getValue() + imm_operand2.getValue();
                        case SUB -> imm_operand1.getValue() - imm_operand2.getValue();
                        case MUL -> imm_operand1.getValue() * imm_operand2.getValue();
                        default -> throw new UnsupportedOperationException("未知的二元操作");
                    };

                    // 使用 MOV 指令将计算结果存储到目标变量
                    Instruction movInstr = Instruction.createMov(target, IRImmediate.of(result));
                    iterator.set(movInstr);
                    continue;
                }
                // 处理立即数作为左操作数的情况
                else if (operand1.isImmediate() && operand2.isIRVariable()) {
                    if (instr_kind == InstructionKind.SUB) {
                        IRVariable tempVar = IRVariable.temp();
                        Instruction movInstr = Instruction.createMov(tempVar, operand1);
                        iterator.previous();
                        iterator.add(movInstr);
                        iterator.next();

                        Instruction subInstr = Instruction.createSub(target, tempVar, operand2);
                        iterator.set(subInstr);
                    } else if (instr_kind == InstructionKind.ADD) {
                        Instruction newInstr = Instruction.createAdd(target, operand2, operand1);
                        iterator.set(newInstr);
                    } else if (instr_kind == InstructionKind.MUL) {
                        IRVariable tempVar = IRVariable.temp();
                        Instruction movInstr = Instruction.createMov(tempVar, operand1);
                        iterator.previous();
                        iterator.add(movInstr);
                        iterator.next();
                        Instruction mulInstr = Instruction.createMul(target, tempVar, operand2);
                        iterator.set(mulInstr);
                    }
                    continue;
                }
                // 处理立即数作为右操作数的乘法情况
                else if (operand2.isImmediate() && operand1.isIRVariable()) {
                    if (instr_kind == InstructionKind.MUL) {
                        IRVariable tempVar = IRVariable.temp();
                        Instruction movInstr = Instruction.createMov(tempVar, operand2);
                        iterator.previous();
                        iterator.add(movInstr);
                        iterator.next();

                        Instruction mulInstr = Instruction.createMul(target, operand1, tempVar);
                        iterator.set(mulInstr);
                    }
                    continue;
                }
            }

            // 如果遇到 return 指令，停止后续指令的处理
            if (instr.getKind().isReturn()) {
                break;
            }
        }
    }

    /**
     * 执行代码生成.
     */
    public void run() {
        RegAllocator regAllocator = new RegAllocator();

        // 添加汇编代码的文本部分
        assembly_instructions.add(".text");

        // 遍历每条中间代码指令并生成对应的汇编指令
        for (Instruction instr : instructions) {

            InstructionKind instr_kind = instr.getKind();
            List<IRValue> operands = instr.getOperands();

            switch (instr_kind) {

                case MOV:
                    IRValue operand = instr.getOperands().getFirst();
                    IRVariable target = instr.getResult();
                    Reg targetReg = regAllocator.allocateRegForVariable(target);

                    // 如果是立即数，使用 `li` 指令加载
                    if (operand instanceof IRImmediate) {
                        assembly_instructions.add("\tli " + targetReg + ", " + ((IRImmediate) operand).getValue());
                    }
                    // 如果是变量，使用 `mv` 指令传递值
                    else if (operand instanceof IRVariable op) {
                        Reg operandReg = regAllocator.allocateRegForVariable(op);
                        assembly_instructions.add("\tmv " + targetReg + ", " + operandReg);

                        // 如果操作数是临时变量，释放对应的寄存器
                        if (op.getName().startsWith("$")) {
                            regAllocator.freeReg(regAllocator.getRegForVariable((IRVariable) operands.getFirst()));
                        }
                    }
                    break;

                case ADD:
                case SUB:
                case MUL:
                    IRVariable dest = instr.getResult();
                    Reg destReg = regAllocator.allocateRegForVariable(dest);
                    Reg op1Reg = regAllocator.allocateRegForVariable((IRVariable) operands.get(0));

                    if (operands.get(1) instanceof IRVariable) {
                        Reg op2Reg = regAllocator.allocateRegForVariable((IRVariable) operands.get(1));
                        String operation = instr_kind == InstructionKind.ADD ? "add" :
                                instr_kind == InstructionKind.SUB ? "sub" : "mul";
                        assembly_instructions.add("\t" + operation + " " + destReg + ", " + op1Reg + ", " + op2Reg);

                        if (((IRVariable) operands.get(0)).getName().startsWith("$")) {
                            regAllocator.freeReg(op1Reg);
                        }
                        if (((IRVariable) operands.get(1)).getName().startsWith("$")) {
                            regAllocator.freeReg(op2Reg);
                        }
                    }
                    // 操作数为立即数，使用 `addi` 或 `subi` 指令
                    else {
                        IRImmediate imm_operand2 = (IRImmediate) operands.get(1);
                        String operation = instr_kind == InstructionKind.ADD ? "addi" : "subi";
                        assembly_instructions.add("\t" + operation + " " + destReg + ", " + op1Reg + ", " + imm_operand2.getValue());
                    }
                    break;

                case RET:
                    IRValue returnVal = instr.getOperands().getFirst();
                    Reg returnReg = regAllocator.getRegForVariable((IRVariable) returnVal);
                    assembly_instructions.add("\tmv a0, " + returnReg);
                    break;

                default:
                    throw new UnsupportedOperationException("未知的指令类型: " + instr_kind);
            }
        }
    }

    /**
     * 输出汇编代码到文件
     *
     * @param path 输出文件路径
     */
    public void dump(String path) {
        FileUtils.writeLines(path, assembly_instructions);
    }
}
